import torch
import numpy as np
import matplotlib.pyplot as plt

# y = 2x

x_data = [1.0, 2.0, 3.0, 4.0, 5.0]

# y = 2x + 1
# y_data = [3.0, 5.0, 7.0, 9.0, 11.0]
y_data = [2.0, 4.0, 6.0, 8.0, 10.0]

w = torch.tensor([1.0])
w.requires_grad_(True)
# b = torch.tensor([0.0])
# b.requires_grad_(True)

'''
    每次做的事情就是先计算出loss，然后求出backward即可，反向传播
'''

def forward(x):
#     return x * w + b
    return w * x


def loss(x, y):
    y_pred = forward(x)
    return (y_pred - y) * (y_pred - y)

loss_list = []
epoch_list = []
print("Predict(before training)", 6, forward(6))

for epoch in range(10):
    for x, y in zip(x_data, y_data):
        l = loss(x, y)
        l.backward()
        print("\tgrad:", x, y, w.grad.item())
        w.data -= w.grad.data * 0.01
        # b.data -= b.grad.data * 0.001

        w.grad.data.zero_()
        # b.grad.data.zero_()
        loss_list.append(l.item())
        epoch_list.append(epoch)

    print("progess:",epoch,l.item())
print("Predict(after training)", 6, forward(6).item())

plt.plot(epoch_list,loss_list)
plt.xlabel("epoch")
plt.ylabel("loss")
plt.show()
